# -*- coding:utf8 -*-
# https://blog.csdn.net/zong596568821xp/article/details/83746615
import cv2
import os
import shutil


def get_frame_from_video(video_name, interval,save_seperate=True,clear_dir=False):
    """
    Args:
        video_name:输入视频名字
        interval: 保存图片的帧率间隔
    Returns:
    """

    # 保存图片的路径
    if save_seperate:
        save_path = video_name.split('.avi')[0] + '/'
    else:
        save_path = '../dataset/MW/'
    save_name_pre = 'Mar'+ video_name.split('.avi')[0].split('-')[-1] +'_'    #
    is_exists = os.path.exists(save_path)
    if not is_exists:
        os.makedirs(save_path)
        print('path of %s is build' % save_path)
    else:
        if clear_dir:
            shutil.rmtree(save_path)
            os.makedirs(save_path)
            print('path of %s already exist and rebuild' % save_path)

    # 开始读视频
    video_capture = cv2.VideoCapture(video_name)
    i = 0
    j = 0

    while True:
        success, frame = video_capture.read()
        i += 1
        if not success:
            print('video is all read')
            break
        if i % interval == 0:
            # 保存图片
            j += 1
            save_name = save_path + save_name_pre + str(i).zfill(6) + '.jpg'
            cv2.imwrite(save_name, frame)
            print('image of %s is saved' % save_name)



if __name__ == '__main__':
    # 视频获取参数设置
    interval = 1                  # 每interval帧采集一幅图像
    video_path = '../dataset'     # 视频保存路径
    video_format = '.avi'
    video_save_seperate = True    # 分开存放视频
    delate_other_file = True      # 删除掉处视频外的其他文件，文件夹中包含其他非视频文件时，需要谨慎选择

    for file in os.listdir(video_path):
        if file.endswith(video_format):
            get_frame_from_video(video_path+'/'+ file, interval,video_save_seperate,video_save_seperate)
        else:
            if delate_other_file:
                shutil.rmtree(video_path+'/'+ file)     # 删除文件夹

    '''
    interval = 1
    video_save_seperate = True
    get_frame_from_video('../dataset/MW-18Mar-3.avi', interval, video_save_seperate)
    '''
